# 自定义的简单线性回归算法
import numpy as np

from MyMLTools.metrics import r2_score


class SimpleLinearRegerssion:
    def __init__(self):
        self.a = None
        self.b = None

    def fit(self, x_train, y_train):
        x_mean = np.mean(x_train)
        y_mean = np.mean(y_train)
        n = (x_train - x_mean).dot(y_train - y_mean)
        d = (x_train - x_mean).dot(x_train - x_mean)
        self.a = n / d
        self.b = y_mean - self.a * x_mean
        return self

    def predict(self, x_predict):
        return np.array([self.__predict(x) for x in x_predict])

    def __predict(self, x):
        return self.a * x + self.b

    def score(self,x_test, y_test):
        y_predict = self.predict(x_test)
        return r2_score(y_test,y_predict)